package jjn.round2;

import commons.TreeNode;

import java.util.ArrayList;
import java.util.HashSet;
import java.util.List;
import java.util.Set;

/**
 * @author Jjn
 * @since 2023/5/30 18:19
 */
public class LeetCode1110_DeleteNodes {
    public List<TreeNode> delNodes(TreeNode root, int[] to_delete) {
        List<TreeNode> result = new ArrayList<>();
        Set<Integer> set = new HashSet<>();
        for (int i : to_delete) {
            set.add(i);
        }
        if (!set.contains(root.val)) {
            result.add(root);
        }
        dfs(root, set, result);
        return result;
    }
    
    private boolean dfs(TreeNode node, Set<Integer> set, List<TreeNode> result) {
        if (node == null) {
            return false;
        }
        if (dfs(node.left, set, result)) {
            node.left = null;
        }
        if (dfs(node.right, set, result)) {
            node.right = null;
        }
        if (set.contains(node.val)) {
            if (node.left != null) {
                result.add(node.left);
            }
            if (node.right != null) {
                result.add(node.right);
            }
            return true;
        }
        return false;
    }
}
